import argparse
import os
import cv2
import sys
import random
import xml.etree.ElementTree as ET

#def process_image(img,min_side):

imgFilePath  = 'niaochao\JPEGImages'
xmlFilePath  = 'niaochao\Annotations'

def cut(fileName,label):
        srcImageFile = os.path.join(imgFilePath,fileName)
        dstImageFile = os.path.join(imgFilePath,label+fileName)
        
        #读image文件
        img = cv2.imread(srcImageFile)
        size = img.shape
        h,w = size[0],size[1]
        x = random.randint(0,w-1664)
        y = random.randint(0,h-1664)
        x2 = x+1664
        y2 = y+1664

        #判断这个区域有没有目标框
        shortName = os.path.splitext(fileName)[0]
        dstXmlFile = os.path.join(xmlFilePath, label+shortName+".xml")
        tree = ET.ElementTree
        tree = ET.parse(os.path.join(xmlFilePath, shortName+".xml"))
        root = tree.getroot()
        #遍历xml中的检测对象，将对象名保存到列表
        objects = []
        for obj in root.findall('object'):
            #判断有没有在裁剪的区域内
            bndbox  = obj.find('bndbox')
            left    = int(bndbox.find('xmin').text)
            top     = int(bndbox.find('ymin').text)
            right   = int(bndbox.find('xmax').text)
            bottom  = int(bndbox.find('ymax').text)
            if ((left >= x)and(top >= y)and(right <= x2)and(bottom <= y2)):
                #修改box位置,保留这个对象
                bndbox.find('xmin').text = str(left - x)
                bndbox.find('ymin').text = str(top - y)
                bndbox.find('xmax').text = str(right - x)
                bndbox.find('ymax').text = str(bottom - y)
                objects.append(obj.find('name').text)
                
            else:
                #删除这个对象           
                root.remove(obj)
        if ((len(objects)>0)and("nest2" in objects)):
            tree.write(dstXmlFile,encoding="utf-8", xml_declaration=False)
            dstimg = img[y:y2,x:x2]
            cv2.imwrite(dstImageFile,dstimg)
            return True
        else:
            return False

if __name__ == '__main__':

    imgFiles = os.listdir(imgFilePath)
    for fileName in imgFiles:
        label = "a"
        for i in range(100):
            if cut(fileName=fileName,label=label):
                print("cut finished:"+label+fileName)
                break
            if (i >= 99):
                print("cut failed:"+label+fileName)
        
        label = "b"
        for i in range(100):            
            if cut(fileName=fileName,label=label):
                print("cut finished:"+label+fileName)
                break
            if (i >= 99):
                print("cut failed:"+label+fileName)

        label = "c"    
        for i in range(100):            
            if cut(fileName=fileName,label=label):
                print("cut finished:"+label+fileName)
                break
            if (i >= 99):
                print("cut failed:"+label+fileName)

        # print(x,x2,y,y2)
        # dstimg = img[y:y2,x:x2]
        # # newimg = cv2.rectangle(img,(0,0),(832,832),(255,255,255),thickness=0)
        # cv2.imwrite(dst,dstimg)


        #长边缩放为min_side
        # scale = min(w,h)/float(832)
        # new_w,new_h = int(w/scale),int(h/scale)
        # resize_img = cv2.resize(img,(new_w,new_h))
        # cv2.imwrite(dst,resize_img)        
        # print(dst)



